# -*- coding: utf-8 -*-
'''
Created on Oct 12, 2010
Decision Tree Source Code for Machine Learning in Action Ch. 3
@author: Peter Harrington & songting
'''
from math import log
import operator


def createDataSet():
    dataSet = [[1, 1, 'yes'],
               [1, 1, 'yes'],
               [1, 0, 'no'],
               [0, 1, 'no'],
               [0, 1, 'no']]
    labels = ['no surfacing', 'flippers']
    # change to discrete values
    return dataSet, labels


def calcShannonEnt(dataSet):
    """
    计算熵, 公式: entropy(S) = -Ｅ Pi * log(Pi, 2)
    :param dataSet: 输入数据集
    :return: 熵
    """
    numEntries = len(dataSet)
    # 分类统计字典
    labelCounts = {}
    ## 统计数据集中各分类
    for featVec in dataSet:  # the the number of unique elements and their occurance
        # 获取当前分类标签
        currentLabel = featVec[-1]
        if currentLabel not in labelCounts.keys(): labelCounts[currentLabel] = 0
        labelCounts[currentLabel] += 1
    shannonEnt = 0.0
    ## 计算所有分类熵的总和
    for key in labelCounts:
        prob = float(labelCounts[key]) / numEntries
        shannonEnt -= prob * log(prob, 2)  # log base 2
    # 返回熵
    return shannonEnt


def splitDataSet(dataSet, axis, value):
    """
    划分数据集
    :param dataSet: 待划分的数据集
    :param axis: 特征的下标
    :param value: 特征值
    :return: 已划分的数据集
    """
    # 创建空的返回数据集
    retDataSet = []
    for featVec in dataSet:
        ## 取出满足特征值的行, 只保留两端的数据, 放入返回的数据集中
        if featVec[axis] == value:
            reducedFeatVec = featVec[:axis]  # chop out axis used for splitting
            reducedFeatVec.extend(featVec[axis + 1:])
            retDataSet.append(reducedFeatVec)
    # 返回已划分的数据集
    return retDataSet


def chooseBestFeatureToSplit(dataSet):
    """
    选取划分数据集最适特征下标
    :param dataSet: 输入的数据集, 类型为 list[list[]]
    :return: 返回最适特征下标
    """
    # 获取特征的数目
    numFeatures = len(dataSet[0]) - 1  # the last column is used for the labels
    # 计算初始熵
    baseEntropy = calcShannonEnt(dataSet)
    # 信息增益
    bestInfoGain = 0.0;
    # 最适特征下标
    bestFeature = -1
    ## 迭代计算, 以每种特征做划分
    for i in range(numFeatures):  # iterate over all the features
        # 取出每列不重复的特征值
        featList = [example[i] for example in dataSet]  # create a list of all the examples of this feature
        uniqueVals = set(featList)  # get a set of unique values
        # 新熵
        newEntropy = 0.0
        ## 迭代计算按当前列划分的新熵
        for value in uniqueVals:
            subDataSet = splitDataSet(dataSet, i, value)
            prob = len(subDataSet) / float(len(dataSet))
            newEntropy += prob * calcShannonEnt(subDataSet)
        # 计算信息增益, 也就是熵的减少量
        infoGain = baseEntropy - newEntropy  # calculate the info gain; ie reduction in entropy
        ## 选取最大的信息增益, 并记录当前特征下标
        if (infoGain > bestInfoGain):  # compare this to the best gain so far
            bestInfoGain = infoGain  # if better than current best, set to best
            bestFeature = i
    # 返回最适特征下标
    return bestFeature  # returns an integer


def majorityCnt(classList):
    """
    投票, 多数表决所属分类
    :param classList: 分类列表
    :return: 返回最多的分类
    """
    classCount = {}
    for vote in classList:
        if vote not in classCount.keys(): classCount[vote] = 0
        classCount[vote] += 1
    sortedClassCount = sorted(classCount.iteritems(), key=operator.itemgetter(1), reverse=True)
    return sortedClassCount[0][0]


def createTree(dataSet, labels):
    """
    创建决策树
    :param dataSet: 数据集
    :param labels: 特征标签列表
    :return: 返回创建的决策树
    """
    subLabels = labels[:]
    # 获取数据集中所有的分类列表
    classList = [example[-1] for example in dataSet]
    ## 所有分类都相同, 返回该分类标签
    if classList.count(classList[0]) == len(classList):
        return classList[0]  # stop splitting when all of the classes are equal
    ## 所有特征已用完, 还未完成分类, 返回最多的分类标签
    if len(dataSet[0]) == 1:  # stop splitting when there are no more features in dataSet
        return majorityCnt(classList)
    # 选择最适特征下标
    bestFeat = chooseBestFeatureToSplit(dataSet)
    # 最适特征标签
    bestFeatLabel = labels[bestFeat]
    # 决策树字典
    myTree = {bestFeatLabel: {}}
    # 删除当前最适特征标签
    del (subLabels[bestFeat])
    ## 获取最适特征所在列的数据, 并去重
    featValues = [example[bestFeat] for example in dataSet]
    uniqueVals = set(featValues)
    ## 为当前列递归创建决策树
    for value in uniqueVals:
        # subLabels = labels[:]  # copy all of labels, so trees don't mess up existing labels
        myTree[bestFeatLabel][value] = createTree(splitDataSet(dataSet, bestFeat, value), subLabels)
    # 返回创建的决策树
    return myTree


def classify(inputTree, featLabels, testVec):
    """
    决策树分类算法
    :param inputTree: 决策树
    :param featLabels: 特征标签列表
    :param testVec: 待分类的向量
    :return: 分类结果
    """
    # 当前决策树的根结点, 特征标签
    firstStr = inputTree.keys()[0]
    # 左右子树
    secondDict = inputTree[firstStr]
    # 取特征标签的值 key
    featIndex = featLabels.index(firstStr)
    key = testVec[featIndex]
    # 根据特征值确定左子树, 还是右子树
    valueOfFeat = secondDict[key]
    ## 如果子树还是字典树, 递归分类
    if isinstance(valueOfFeat, dict):
        classLabel = classify(valueOfFeat, featLabels, testVec)
    ## 如果到达叶子结点, 返回该结点的分类结果
    else:
        classLabel = valueOfFeat
    # 返回分类结果
    return classLabel


def storeTree(inputTree, filename):
    """
    序列化决策树
    :param inputTree: 输入的决策树
    :param filename: 存储位置
    :return:
    """
    import pickle
    fw = open(filename, 'w')
    pickle.dump(inputTree, fw)
    fw.close()


def grabTree(filename):
    """
    反序列化决策树
    :param filename: 存储位置
    :return: 反序列化忽的决策树
    """
    import pickle
    fr = open(filename)
    return pickle.load(fr)
